package bma.common.thrift;

import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.util.ArrayList;
import java.util.List;

import org.apache.thrift.transport.TNonblockingSocket;
import org.apache.thrift.transport.TServerTransportAcceptFilter;
import org.apache.thrift.transport.TSocket;
import org.apache.thrift.transport.TTransport;

import bma.common.langutil.io.InetNetwork;

public class SecureTServerTransportAcceptFilter implements
		TServerTransportAcceptFilter {

	final org.slf4j.Logger log = org.slf4j.LoggerFactory
			.getLogger(SecureTServerTransportAcceptFilter.class);

	private List<InetNetwork> whiteList = new ArrayList<InetNetwork>();
	private List<InetNetwork> blackList = new ArrayList<InetNetwork>();

	public void setWhiteList(List<String> whiteList) {
		for (String s : whiteList) {
			this.whiteList.add(InetNetwork.create(s));
		}
	}

	public void setWhiteAddress(String whiteList) {
		if (whiteList == null || whiteList.isEmpty()) {
			return;
		}

		String[] ips = whiteList.split(",");
		for (String ip : ips) {
			this.whiteList.add(InetNetwork.create(ip.trim()));
		}
	}

	public void setBlackList(List<String> list) {
		for (String s : list) {
			this.blackList.add(InetNetwork.create(s));
		}
	}

	public void setBlackAddress(String blackList) {
		if (blackList == null || blackList.isEmpty()) {
			return;
		}

		String[] ips = blackList.split(",");
		for (String ip : ips) {
			this.blackList.add(InetNetwork.create(ip.trim()));
		}
	}

	@Override
	public TTransport filter(TTransport transport) {
		SocketAddress addr = null;
		if (transport instanceof TSocket) {
			TSocket s = (TSocket) transport;
			addr = s.getSocket().getRemoteSocketAddress();
		}
		if (transport instanceof TNonblockingSocket) {
			TNonblockingSocket s = (TNonblockingSocket) transport;
			addr = s.getSocketChannel().socket().getRemoteSocketAddress();
		}
		if (addr == null) {
			if (log.isWarnEnabled()) {
				log.warn("transport '" + transport.getClass().getName()
						+ "' not supported");
			}
			return null;
		}
		if (checkAddress(addr)) {
			return transport;
		}
		return null;
	}

	public boolean checkAddress(SocketAddress addr) {
		if (addr instanceof InetSocketAddress) {
			InetAddress iaddr = ((InetSocketAddress) addr).getAddress();
			CheckResult cr = getAddressResult(iaddr);
			if (cr.equals(CheckResult.IN_BLACKLIST)) {
				if (log.isWarnEnabled()) {
					log.warn("blacklist address({})", iaddr);
				}
				return false;
			}
			if (cr.equals(CheckResult.NOT_IN_WHITE_LIST)) {
				if (log.isWarnEnabled()) {
					log.warn("not whitelist address({})", iaddr);
				}
				return false;
			}
		}
		if (log.isDebugEnabled()) {
			log.debug("accept address({})", addr);
		}
		return true;
	}

	public enum CheckResult {
		ALL_PASS, IN_BLACKLIST, NOT_IN_BLACKLIST, NOT_IN_WHITE_LIST,
	}

	public CheckResult getAddressResult(InetAddress iaddr) {
		if (!blackList.isEmpty()) {
			for (InetNetwork n : blackList) {
				if (n.isMatch(iaddr)) {
					return CheckResult.IN_BLACKLIST;
				}
			}
		}
		if (!whiteList.isEmpty()) {
			boolean match = false;
			for (InetNetwork n : whiteList) {
				if (n.isMatch(iaddr)) {
					match = true;
					break;
				}
			}
			if (!match) {
				return CheckResult.NOT_IN_WHITE_LIST;
			}
			return CheckResult.ALL_PASS;
		}
		return CheckResult.NOT_IN_BLACKLIST;
	}
}
